import torch.nn as nn 
import torch.nn.functional as F 

class CNN(nn.Module): 
    def __init__(self): 
        super(CNN,self).__init__()

        self.conv_net = nn.Sequential( 
            # Conv Block 1
            nn.Conv2d(in_channels=1, out_channels=20, kernel_size=5, stride=1), 
            nn.ReLU(inplace=True), 
            nn.MaxPool2d(kernel_size=2, stride=2),

            # Conv Block 2 
            nn.Conv2d(in_channels=20, out_channels=50, kernel_size=5, stride=1),
            nn.MaxPool2d(kernel_size=2, stride=2), 
            nn.ReLU(inplace=True)
        )

        self.fc = nn.Sequential( 
            nn.Linear(in_features=4*4*50, out_features=500), 
            nn.ReLU(inplace=True),
            nn.Linear(in_features=500, out_features=10)
        )   

    def forward(self, x): 
        x = self.conv_net(x) 
        x = x.view(-1, 4*4*50)
        x = self.fc(x) 
        return x